import asyncio
import websockets
from loguru import logger
from typing import Optional, Callable, Any, Dict
from ..handlers.message_handler import DouyinMessageHandler
from ...utils.proxy_manager import GlobalProxyManager
from python_socks.async_.asyncio import Proxy
from ..config.douyin_config import DouyinConfig


class WebSocketConnectionManager:
    def __init__(self,
                 message_handler: DouyinMessageHandler,
                 proxy_manager: GlobalProxyManager) -> None:
        """
        初始化WebSocket连接管理器

        Args:
            message_handler: 消息处理器实例
            proxy_manager: 代理管理器实例
        """
        self.message_handler = message_handler
        self.proxy_manager = proxy_manager
        self.active_connections: Dict[str,
                                      websockets.WebSocketClientProtocol] = {}

    async def create_connection(self,
                                wss_url: str,
                                ttwid: str,
                                web_url: str,
                                shared_queue: Any,
                                shared_list: list,
                                on_close_callback: Callable) -> None:
        """创建异步WebSocket连接"""
        max_retries = 3
        retry_count = 0
        
        while retry_count < max_retries:
            try:
                proxy = self.proxy_manager.get_proxy()
                proxy_host, proxy_port = proxy.split(':')
                proxy = Proxy.from_url(f'socks5://{proxy_host}:{proxy_port}')

                headers = {
                    'User-Agent': "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/96.0.4664.55 Safari/537.36",
                    'Cookie': f"ttwid={ttwid}"
                }

                host = wss_url.split('/')[2]

                async with websockets.connect(
                    wss_url,
                    additional_headers=headers,
                    ping_interval=21,
                    ping_timeout=20,
                    close_timeout=10,
                    sock=await proxy.connect(
                        dest_host=host,
                        dest_port=443
                    )
                ) as websocket:
                    self.active_connections[web_url] = websocket
                    # logger.info(f"WebSocket连接成功: {web_url}")

                    try:
                        async for message in websocket:
                            await self._handle_message(
                                message, web_url, websocket,
                                shared_queue, shared_list, on_close_callback
                            )

                    except websockets.exceptions.ConnectionClosed:
                        # logger.warning(f"WebSocket连接关闭: {web_url}")
                        await self._handle_connection_closed(
                            web_url, shared_list, on_close_callback
                        )
                    return  # 连接成功，退出重试循环

            except Exception as e:
                retry_count += 1
                if retry_count < max_retries:
                    # logger.error(f"WebSocket连接失败: {web_url} - {str(e)}，准备第 {retry_count + 1} 次重试")
                    await asyncio.sleep(2)  # 重试前等待
                    # 尝试获取新代理
                    if "timed out" in str(e):
                        self.proxy_manager.get_new_proxy()
                else:
                    # logger.error(f"WebSocket连接失败: {web_url} - {str(e)}，已达到最大重试次数")
                    await self._handle_connection_closed(web_url, shared_list, on_close_callback)

    async def _handle_message(self, message: str, web_url: str,
                              websocket: websockets.WebSocketClientProtocol,
                              shared_queue: Any, shared_list: list,
                              on_close_callback: Callable):
        """处理接收到的消息"""
        try:
            await self.message_handler.handle_message(
                message, web_url, websocket,
                shared_queue, shared_list, on_close_callback
            )
        except Exception as e:
            logger.error(f"处理消息错误: {web_url} - {str(e)}")

    async def _handle_connection_closed(self, web_url: str,
                                        shared_list: list,
                                        on_close_callback: Callable):
        """处理连接关闭"""
        if web_url in self.active_connections:
            del self.active_connections[web_url]
        if on_close_callback:
            await on_close_callback(None, web_url, shared_list, "连接关闭")

    async def _handle_error(self, web_url: str, error_msg: str) -> bool:
        """处理错误情况"""
        if "454" in error_msg:
            self.proxy_manager.get_new_proxy()
            return True

        if any(err in error_msg for err in [
            "failed CONNECT via proxy",
            "WinError",
            "EOF occurred in violation of protocol",
            "connect call failed",
            "connection closed abnormally"
            'unreachable'
        ]):
            logger.info(f"{web_url} 错误: {error_msg}")
            return True

        return False

    async def close_all_connections(self):
        """关闭所有活跃的连接"""
        for url, ws in self.active_connections.items():
            try:
                await ws.close()
                logger.info(f"关闭WebSocket连接: {url}")
            except Exception as e:
                logger.error(f"关闭WebSocket连接失败: {url} - {str(e)}")
        self.active_connections.clear()
